package com.ruben.simplescaffold.filter.wrappers;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.TypeReference;
import com.ruben.simplescaffold.utils.sensitive.SensitiveWordUtils;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.ServletRequest;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.*;

/**
 * 封装请求
 *
 * @author <achao1441470436@gmail.com>
 * @since 2021/8/7 21:50
 */
public class RequestWrapper extends HttpServletRequestWrapper {

    private final Map<String, String[]> parameterMap;
    private final byte[] body;

    public RequestWrapper(HttpServletRequest request, SensitiveWordUtils sensitiveWordUtils, Map<String, String[]> parameterMap, String bodyString) {
        super(request);
        if (sensitiveWordUtils == null) {
            this.parameterMap = Optional.ofNullable(parameterMap).orElseGet(super::getParameterMap);
            queryStringPutParameterMap(super.getQueryString());
            this.body = Optional.ofNullable(bodyString).orElseGet(() -> BodyHelper.getBodyString(request)).getBytes(StandardCharsets.UTF_8);
        } else {
            String paramString = JSON.toJSONString(Optional.ofNullable(parameterMap).orElseGet(super::getParameterMap));
            paramString = sensitiveWordUtils.replaceSensitiveWord(paramString, SensitiveWordUtils.MAX_MATCH_TYPE, SensitiveWordUtils.REPLACE_CHAR);
            this.parameterMap = JSON.parseObject(paramString, new TypeReference<Map<String, String[]>>() {
            });
            Optional.ofNullable(super.getQueryString()).ifPresent(queryString -> queryStringPutParameterMap(sensitiveWordUtils.replaceSensitiveWord(queryString, SensitiveWordUtils.MAX_MATCH_TYPE, SensitiveWordUtils.REPLACE_CHAR)));
            this.body = sensitiveWordUtils.replaceSensitiveWord(Optional.ofNullable(bodyString).orElseGet(() -> BodyHelper.getBodyString(request)), SensitiveWordUtils.MAX_MATCH_TYPE, SensitiveWordUtils.REPLACE_CHAR).getBytes(StandardCharsets.UTF_8);
        }

    }

    @Override
    public String getParameter(String name) {
        String[] values = parameterMap.get(name);
        if (Objects.isNull(values) || values.length == 0) {
            return null;
        }
        return values[0];
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        if (parameterMap == null) {
            return super.getParameterMap();
        }
        return parameterMap;
    }

    @Override
    public Enumeration<String> getParameterNames() {
        if (parameterMap == null) {
            return super.getParameterNames();
        }
        return new Vector<>(parameterMap.keySet()).elements();
    }

    @Override
    public String[] getParameterValues(String name) {
        if (parameterMap == null) {
            return super.getParameterValues(name);
        }
        return parameterMap.get(name);
    }

    @Override
    public BufferedReader getReader() {
        return new BufferedReader(new InputStreamReader(getInputStream()));
    }

    @Override
    public ServletInputStream getInputStream() {
        final ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(body);
        return new ServletInputStream() {
            @Override
            public int read() {
                return byteArrayInputStream.read();
            }

            @Override
            public boolean isFinished() {
                return false;
            }

            @Override
            public boolean isReady() {
                return false;
            }

            @Override
            public void setReadListener(ReadListener arg0) {
            }
        };
    }

    @Override
    public String getHeader(String name) {
        return super.getHeader(name);
    }

    @Override
    public Enumeration<String> getHeaderNames() {
        return super.getHeaderNames();
    }

    @Override
    public Enumeration<String> getHeaders(String name) {
        return super.getHeaders(name);
    }

    private void queryStringPutParameterMap(String queryString) {
        if (queryString != null && queryString.trim().length() > 0) {
            String[] params = queryString.split("&");
            for (String param : params) {
                int splitIndex = param.indexOf("=");
                if (splitIndex == -1) {
                    continue;
                }
                String key = param.substring(0, splitIndex);
                if (!this.parameterMap.containsKey(key)) {
                    if (splitIndex < param.length()) {
                        String value = param.substring(splitIndex + 1);
                        this.parameterMap.put(key, new String[]{value});
                    }
                }
            }
        }
    }


    public static class BodyHelper {
        /**
         * 获取请求Body
         *
         * @param request
         * @return
         */
        public static String getBodyString(ServletRequest request) {
            StringBuilder sb = new StringBuilder();
            InputStream inputStream = null;
            BufferedReader reader = null;
            try {
                inputStream = request.getInputStream();
                reader = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8));
                String line = "";
                while ((line = reader.readLine()) != null) {
                    sb.append(line);
                }
            } catch (IOException e) {
                e.printStackTrace();
            } finally {
                if (inputStream != null) {
                    try {
                        inputStream.close();
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                }
                if (reader != null) {
                    try {
                        reader.close();
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                }
            }
            return sb.toString();
        }
    }
}
